from swarm import Swarm, Result
import json
import json
import os
from datetime import datetime, timedelta
from queue import Queue
from time import sleep
import speech_recognition as sr
import torch
import numpy as np
import whisper
import argparse
from queue import Empty
def process_and_print_streaming_response(response):
    content = ""
    last_sender = ""
    for chunk in response:
        if "sender" in chunk:
            last_sender = chunk["sender"]
        if "content" in chunk and chunk["content"] is not None:
            if not content and last_sender:
                print(f"\033[94m{last_sender}:\033[0m", end=" ", flush=True)
                last_sender = ""
            print(chunk["content"], end="", flush=True)
            content += chunk["content"]
        if "tool_calls" in chunk and chunk["tool_calls"] is not None:
            for tool_call in chunk["tool_calls"]:
                f = tool_call["function"]
                name = f["name"]
                if not name:
                    continue
                print(f"\033[94m{last_sender}: \033[95m{name}\033[0m()")
        if "delim" in chunk and chunk["delim"] == "end" and content:
            print()
            content = ""
        if "response" in chunk:
            return chunk["response"]


def pretty_print_messages(messages) -> None:
    for message in messages:
        if message["role"]!= "assistant":
            continue
        print(f"\033[94m{message['sender']}\033[0m:", end=" ")
        if message["content"]:
            print(message["content"])
        tool_calls = message.get("tool_calls") or []
        if len(tool_calls) > 1:
            print()
        for tool_call in tool_calls:
            f = tool_call["function"]
            name, args = f["name"], f["arguments"]
            arg_str = json.dumps(json.loads(args)).replace(":", "=")
            print(f"\033[95m{name}\033[0m({arg_str[1:-1]})")

def pretty_format_messages(messages) -> list:
    """
    将消息列表进行格式化处理，返回格式化后的消息内容列表，方便后续传递使用。

    参数:
    messages (list): 包含消息字典的列表，每个字典包含角色、发送者、内容、工具调用等信息。

    返回:
    list: 格式化后的消息内容字符串列表。
    """
    formatted_messages = []
    for message in messages:
        if message["role"]!= "assistant":
            continue
        msg_str = ""
        msg_str += f"{message['sender']}: "
        if message["content"]:
            msg_str += message["content"]
        tool_calls = message.get("tool_calls") or []
        if len(tool_calls) > 1:
            msg_str += "\n"
        for tool_call in tool_calls:
            f = tool_call["function"]
            name, args = f["name"], f["arguments"]
            arg_str = json.dumps(json.loads(args)).replace(":", "=")
            msg_str += f"{name}({arg_str[1:-1]})"
        formatted_messages.append(msg_str)
    return formatted_messages

def run_demo_loop(
        openai_client,
        starting_agent,
        context_variables=None,
        stream=False,
        debug=False
) -> None:
    client = Swarm(openai_client)
    print("Starting Swarm CLI 🐝")
    print('Type "exit" or "quit" to leave the chat.')

    messages = []
    agent = starting_agent

    while True:
        user_input = input("\033[90mUser\033[0m: ").strip()
        if user_input.lower() in {"exit", "quit"}:
            with open('knowledgeBase/Keyword statistics.json', 'w', encoding='UTF-8') as f:
                json.dump(context_variables, f, ensure_ascii=False)
            print("Exiting chat. Goodbye!")
            break
        messages.append({"role": "user", "content": user_input})

        response = client.run(
            agent=agent,
            messages=messages,
            context_variables=context_variables or {},
            stream=stream,
            debug=debug,
            max_turns=100
        )

        if stream:
            response = process_and_print_streaming_response(response)
        else:
            pretty_print_messages(response.messages)

        messages.extend(response.messages)
        agent = response.agent
        context_variables = response.context_variables

def run_demo_loop_with_voice(
        openai_client,
        starting_agent,
        context_variables=None,
        stream=False,
        debug=False
) -> None:
    client = Swarm(openai_client)
    print("Starting Swarm CLI 🐝")
    print('Type "exit" or "quit" to leave the chat.')

    messages = []
    agent = starting_agent

    # 用于语音转文字的相关初始化，参考之前语音转文字代码的逻辑
    phrase_time = None
    data_queue = Queue()
    recorder = sr.Recognizer()
    recorder.energy_threshold = 1000
    recorder.dynamic_energy_threshold = False
    source = sr.Microphone(sample_rate=16000)
    with source:
        recorder.adjust_for_ambient_noise(source)

    def record_callback(_, audio: sr.AudioData) -> None:
        """
        线程化的回调函数，用于在录音完成时接收音频数据
        audio: 包含录制字节数据的AudioData对象
        """
        data = audio.get_raw_data()
        data_queue.put(data)

    recorder.listen_in_background(source, record_callback, phrase_time_limit=2)

    while True:
        try:
            if not data_queue.empty():
                now = datetime.utcnow()
                phrase_complete = False
                if phrase_time and now - phrase_time > timedelta(seconds=3):
                    phrase_complete = True
                phrase_time = now

                audio_data = b''.join(data_queue.queue)
                data_queue.queue.clear()
                audio_np = np.frombuffer(audio_data, dtype=np.int16).astype(np.float32) / 32768.0
                # 这里加载whisper模型，假设采用之前语音转文字代码中的模型加载逻辑（比如默认模型等）
                model = "base"
                audio_model = whisper.load_model(model)
                result = audio_model.transcribe(audio_np, fp16=torch.cuda.is_available())
                user_input = result['text'].strip()

                if user_input.lower() in {"exit", "quit"}:
                    with open('knowledgeBase/Keyword statistics.json', 'w', encoding='UTF-8') as f:
                        json.dump(context_variables, f, ensure_ascii=False)
                    print("Exiting chat. Goodbye!")
                    break
                messages.append({"role": "user", "content": user_input})

                response = client.run(
                    agent=agent,
                    messages=messages,
                    context_variables=context_variables or {},
                    stream=stream,
                    debug=debug,
                    max_turns=100
                )

                if stream:
                    response = process_and_print_streaming_response(response)
                else:
                    pretty_print_messages(response.messages)

                messages.extend(response.messages)
                agent = response.agent
                context_variables = response.context_variables
            else:
                # 若语音队列空，适当休眠避免过度占用CPU
                sleep(0.25)
        except KeyboardInterrupt:
            break

def run_demo_loop_to_backend(
        openai_client,
        starting_agent,
        context_variables=None,
        input_queue=None,
        socketio_instance=None,
        stream=False,
        debug=False,
         # 添加一个参数用于接收输入队列
) -> None:
    client = Swarm(openai_client)
    print("Starting Swarm CLI 🐝")
    print('Type "exit" or "quit" to leave the chat.')

    messages = []
    agent = starting_agent

    if input_queue is None:
        raise ValueError("input_queue parameter cannot be None")

    while True:
        try:
            user_input = input_queue.get(block=True, timeout=1)  # 从队列中获取输入内容，设置超时时间
            user_input = user_input.strip()
            if user_input.lower() in {"exit", "quit"}:
                with open('knowledgeBase/Keyword statistics.json', 'w', encoding='UTF-8') as f:
                    json.dump(context_variables, f, ensure_ascii=False)
                print("Exiting chat. Goodbye!")
                break
            messages.append({"role": "user", "content": user_input})

            response = client.run(
                agent=agent,
                messages=messages,
                context_variables=context_variables or {},
                stream=stream,
                debug=debug,
                max_turns=100
            )

            if stream:
                response = process_and_print_streaming_response(response)
            else:
                pretty_print_messages(response.messages)
                formatted_messages = pretty_format_messages(response.messages)
                for formatted_message in formatted_messages:
                    socketio_instance.emit('new_message', {'message': formatted_message})

            messages.extend(response.messages)
            agent = response.agent
            context_variables = response.context_variables
        except Empty:
            continue

def run_demo_loop_to_backend1(
        openai_client,
        starting_agent,
        context_variables=None,
        input_queue=None,
        output_queue=None,
        socketio_instance=None,
        stream=False,
        debug=False,
         # 添加一个参数用于接收输入队列
) -> None:
    client = Swarm(openai_client)
    print("Starting Swarm CLI 🐝")
    print('Type "exit" or "quit" to leave the chat.')

    messages = []
    agent = starting_agent

    if input_queue is None:
        raise ValueError("input_queue parameter cannot be None")

    while True:
        try:
            user_input = input_queue.get(block=True, timeout=1)  # 从队列中获取输入内容，设置超时时间
            user_input = user_input.strip()
            if user_input.lower() in {"exit", "quit"}:
                with open('knowledgeBase/Keyword statistics.json', 'w', encoding='UTF-8') as f:
                    json.dump(context_variables, f, ensure_ascii=False)
                print("Exiting chat. Goodbye!")
                break
            messages.append({"role": "user", "content": user_input})

            response = client.run(
                agent=agent,
                messages=messages,
                context_variables=context_variables or {},
                stream=stream,
                debug=debug,
                max_turns=100
            )

            if stream:
                response = process_and_print_streaming_response(response)
            else:
                pretty_print_messages(response.messages)
                formatted_messages = pretty_format_messages(response.messages)
                output_queue.put(formatted_messages)
                #for formatted_message in formatted_messages:
                 #   socketio_instance.emit('new_message', {'message': formatted_message})

            messages.extend(response.messages)
            agent = response.agent
            context_variables = response.context_variables
        except Empty:
            continue